//////////////////////////////////////////////////////////////////
//                        COPYRIGHT NOTICE
// (c) Copyright 1998-2002 by Jeongnim Kim, Lars Jonsson and Seungwon Lee
//   All rights reserved.
//////////////////////////////////////////////////////////////////

//////////////////////////////////////////////////////////////////
//                    CONTACT INFORMATION
//   Jeongnim Kim
//   National Center for Supercomputing Applications &
//   Materials Computation Center
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//   Tel:    217-244-6319 (NCSA) 217-333-3324 (MCC)
//
// Supported by 
//   National Center for Supercomputing Applications, UIUC
//   Materials Computation Center, UIUC
//   Department of Physics, Ohio State University
//   Ohio Supercomputer Center
//////////////////////////////////////////////////////////////////
// -*- C++ -*-

#ifndef OHMMS_OMP_ARPACKSYMENGINE_H
#define OHMMS_OMP_ARPACKSYMENGINE_H

#include <complex>
#include <iostream>
#include <map>
#include <algorithm>
using namespace std;
#include <omp.h>

#include "EigSolver/EigenSpaceDimension.h"
#include "EigSolver/Arpackf.h"
#include "OhmmsNumeric/OhmmsBlasf.h"
#ifdef OHMMS_MTL
#include "mtl/mtl.h"
#else
#include "OhmmsPETE/OhmmsVector.h"
#endif

#if defined(ADD_)
#define dsaupd_mp dsaupd_mp_
#define znaupd_mp znaupd_mp_
#endif

extern "C" {

  void dsaupd_mp(int* idmy,
              int *ido, const char *bmat, int *n, const char *which,
              int *nev, double *tol, double *resid,
              int *ncv, double *V, int *ldv,
              int *iparam, int *ipntr, double *workd,
              double *workl, int *lworkl, int *info);

  void znaupd_mp(int* idmy,
              int *ido, const char* bmat, int *n, const char* which,
              int *nev, double *tol, complex<double> *resid,
              int *ncv, complex<double> *V, int *ldv,
              int *iparam, int *ipntr, complex<double> *workd,
              complex<double> *workl, int *lworkl, double *rwork,
              int *info); 

}


template<class T, class EigVec_t>
struct ArpackSymSolverOMP {

  typedef typename OHMMS::type_traits<T>::value_type value_type;
  typedef typename OHMMS::type_traits<T>::real_type  real_type;
  typedef EigVec_t                 EigenStates_t;
#ifdef OHMMS_MTL
  typedef mtl::dense1D<value_type> WorkArray_t;
  typedef mtl::dense1D<real_type>  RealArray_t;
  typedef mtl::dense1D<int>        IndexArray_t;
#else
  typedef Vector<value_type> WorkArray_t;
  typedef Vector<real_type>  RealArray_t;
  typedef Vector<int>        IndexArray_t;
#endif
  typedef typename WorkArray_t::size_type   size_type;

public:

  EigenSpaceDimension& eigStates;
  int N;
  int M;
  int ncv; // number of Lanczos vectors
  int info;
  int lworkl;
  int ido;
  int ierr;
  int nconv;
  int maxiter;
  int ishfts;
  int mode;
  int rvec;
  char bmat;
  string which;
  int iparam[11], ipntr[14];

  IndexArray_t select;
  WorkArray_t workd; // array of 3N+1
  WorkArray_t workl; // array of lworkl
  WorkArray_t resid; // residue
  WorkArray_t eigVal;// eigen values
  WorkArray_t eigVec;// eigen vectors 
  RealArray_t rwork; // only for a complex problem

  real_type tol;  // tolerance
  value_type sigma; // reference value for shift-invert

  //OMP specifics
  int numthreads, maxmpblock;
  IndexArray_t mpdist,atdist;
  
  ArpackSymSolverOMP(EigenSpaceDimension& dim, bool nonortho=false):
    eigStates(dim) {   
    tol = 1.0e-13;
    info = 0;
    ido = 0;
    maxiter = 1000;
    ncv = -1;
    numthreads = 0;

    which = "SM";
    iparam[2] = maxiter;
  }

  ~ArpackSymSolverOMP() { }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H) {

    bmat = 'I';
    initializeWorkSpace(H);
    iparam[0] = 1; iparam[6] = 1;
    ido = 0;

    do{
      ///arpack
      solve(eigVec.data());

      ///multiplication
      H.apply(workd.data()+ipntr[0]-1, workd.data()+ipntr[1]-1);

    } while(ido == -1 || ido == 1);
    rvec = 1;
    finalize(eigVec.data());
    if(info < 0) {
      cout << endl << "Error with __aupd, info = " <<  info << endl;
      cout << " Check documentation in __aupd " << endl << endl;
      return false;
    } else
      return true;
  }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H, EigenStates_t& Psi) {
    bool done = solve(H);
    // copy eigen states according to the ascending eigenvalues
    multimap<double,int> ordered_e;
    for(int i=0; i<M; i++) ordered_e.insert(pair<double,int>(eigVal[i],i));
    multimap<double,int>::iterator it = ordered_e.begin();
    int ieig = 0;
    for(; it!=ordered_e.end(); it++) {
      eigVal[ieig] = (*it).first;      
      int is = (*it).second*N;
      std::copy(eigVec.data()+is,eigVec.data()+is+N, Psi.data()+ieig*N);
      ieig++; 
    }
    return done;
  }


  template<class HamiltonianMatrix_t>
  bool getEigvalues(HamiltonianMatrix_t& H) {

    initializeWorkSpace(H);
    do{
      solve(eigVec.data());
      H.apply(workd.data()+ipntr[0]-1, workd.data()+ipntr[1]-1);
    } while(ido == -1 || ido == 1);

    rvec = 0;

    finalize(eigVec.data());

    if(info < 0) {
      cout << endl << "Error with __aupd, info = " <<  info << endl;
      cout << " Check documentation in __aupd " << endl << endl;
      return false;
    } else
      return true;
  }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H, HamiltonianMatrix_t& S) { 
    ///not implemented yet
    cerr << "ArpackSymSolverOMP does not work with a generalized eigen problems"
	 << endl;
    return false;
  }

  template<class HamiltonianMatrix_t>
  bool solve(HamiltonianMatrix_t& H, HamiltonianMatrix_t& S, 
	     EigenStates_t& Psi) {
    ///not implemented yet
    cerr << "ArpackSymSolverOMP does not work with a generalized eigen problems"
	 << endl;
    return false;
  }

  template<class HamiltonianMatrix_t>
  bool checkEigvectors(HamiltonianMatrix_t& H, ostream& os) {
    if ( ierr != 0) {
      os << endl << "Error with __eupd, info = " << ierr << endl;
      os << "Check the documentation of __eupd. " << endl;
      return false;
    } else {
      
      nconv =  iparam[4];
      if(rvec == 1) {
	WorkArray_t ax(N);
	WorkArray_t res(M);
	
	for(int j=0; j<nconv; j++) { // residual calculations
	  H.apply(eigVec.data()+j*N, ax.data());
	  value_type da  = -eigVal[j];
	  int incx = 1;
	  int incy = 1;
	  OhmmsBlas::axpy(N, da, eigVec.data()+j*N, incx, ax.data(), incy);  
	  res[j] = OhmmsBlas::norm2(N, ax.data(), incx);
	  res[j] /= fabs(eigVal[j]);
	}
	os << "Ritz values and relative residuals " << endl;
	for(int j=0; j<nconv; j++)
	  os << j << "  "  << eigVal[j] << "  " << res[j] << endl;
      } else {
	os << "Ritz values only " << endl;
	for(int j=0; j<nconv; j++)
	  os << j << "  "  << eigVal[j] <<  endl;
      }

      os << endl;
      if(info == 1) {
	os << endl << " Maximum number of iterations reached." << endl;
      } else if ( info == 3) {
	os << endl << "No shifts could be applied during implicit "
	     <<   "  Arnoldi update, try increasing NCV." << endl;
      }
      os << endl;
      os << " _SDRV1 " << endl;
      os << " ====== " << endl << endl;
      os << " Size of the matrix is " << N << endl;
      os << " The number of Ritz values requested is " << M << endl;
      os << " The number of Arnoldi vectors generated  (NCV) is " << ncv << endl;
      os << " What portion of the spectrum: " << which << endl;
      os << " The number of converged Ritz values is " << nconv << endl;
      
      os << " The number of Implicit Arnoldi update iterations taken is ";
      os <<  iparam[2] << endl;
      os << " The number of OP*x is " << iparam[8] << endl;
      os << " The convergence criterion is " << tol << endl;
      return true;
    }
  }

  template<class HamiltonianMatrix_t>
  bool checkEigvectors(HamiltonianMatrix_t& H, HamiltonianMatrix_t& S,
		       ostream& os) {
    return false;
  }
  inline void setMaxIter(int n) { maxiter = n;}
  inline void setMaxTrialStates(int n) { ncv = n;}
  inline void setTolerance(real_type x) { tol = x;}
  inline real_type eigval(size_type i) const { return eigVal[i];}

//   bool put(xmlDocPtr doc, xmlNsPtr ns, xmlNodePtr cur);

  template<class HamiltonianMatrix_t>
  void initializeWorkSpace(HamiltonianMatrix_t& H) {
//     // getting dimension from EigenSpace
//     N = States.N;
//     M = States.M;

    if(ncv < 0) ncv = 4*M;
    if(ncv > N) ncv = N;
    //TODO: COMPLEX and DOUBLE  (maybe not)
    lworkl = ncv*(3*ncv+5);
    iparam[0] = 1; iparam[6] = 1;
    if(workd.size() < 3*N) {
      workd.resize(3*N);
    }                                    
    if(resid.size() < N) {
      resid.resize(N);
    }
    if(select.size() < ncv) {
      select.resize(ncv);
    }
    if(eigVal.size() < ncv*2) {
      eigVal.resize(ncv*2);
    }
    if(eigVec.size() < ncv*N) {
      eigVec.resize(ncv*N);
    }
    if(workl.size() < lworkl) {
      workl.resize(lworkl);
    }
    
    if(mpdist.size() == 0) {
#ifdef ENABLE_OPENMP
      numthreads = omp_get_max_threads();
#else
      numthreads = 1;
#endif
      mpdist.resize(numthreads+1);
      atdist.resize(numthreads+1);

      atdist[0] = 0; // block-based distribution
      mpdist[0] = 0; // element-based distribution

      int nloc = N/numthreads;
      int dn = N- nloc*numthreads;
      maxmpblock = H.numcol(0);

      for(int myid = 1; myid < numthreads; myid++) {
	atdist[myid] = H.numcol(myid-1) + atdist[myid-1];
	maxmpblock = max(maxmpblock, atdist[myid]-atdist[myid-1]);
	if(myid <= dn) mpdist[myid] = mpdist[myid-1]+nloc + 1;
	else           mpdist[myid] = mpdist[myid-1]+nloc;
      }
      atdist[numthreads] = N;
      mpdist[numthreads] = N;

      maxmpblock = max(maxmpblock,N-atdist[numthreads-1]);

      cout << "Vector distribution: block-based element-base\n";
      for(int myid = 0; myid <= numthreads; myid++) {
	cout << "                " 
	     << atdist[myid] << " " << mpdist[myid] << endl;
      }

      if(dn != 0) nloc++;
      if(maxmpblock < nloc) {
	cout << "nloc " << nloc << endl;
	cout << "maxmpblock " << maxmpblock << endl;
	cout << "Something is very wrong. The local block based on "
	     << " block matrix MUST be larger.\n";
      }
    }
  }


  void solve(double* eigv){
    int n = N; int ldv = N; int nev = M;
    dsaupd_mp(mpdist.data(),
	      &ido, &bmat, &n, which.c_str(), &nev, &tol,
	      resid.data(), &ncv, eigv, &ldv, iparam, ipntr,
	      workd.data(), workl.data(), &lworkl, &info);
  }

  void solve(complex<double>* eigv){

    if(rwork.empty()) {
      rwork.resize(ncv); // only for a complex problem
    } 
    int n = N; int ldv = N; int nev = M;
    znaupd_mp(mpdist.data(),
	      &ido, &bmat, &n, which.c_str(), &nev, &tol,
	      resid.data(), &ncv, eigv, &ldv, iparam, ipntr,
	      workd.data(), workl.data(), &lworkl, rwork.data(), &info);
  }

  void finalize(double *eigv) {
    char howmany ='A';
    int n = N; int ldv = N; int nev = M;
    dseupd( &rvec, &howmany, select.data(),
            eigVal.data(), eigv, &ldv, &sigma,
            &bmat, &n, which.c_str(), &nev, &tol, resid.data(), &ncv,
            eigv, &ldv, iparam, ipntr,
            workd.data(), workl.data(), &lworkl, &ierr );

  }

  void finalize(complex<double> *eigv) {
    char howmany ='A';
    int n = N; int ldv = N; int nev = M;
    WorkArray_t workev(2*ncv);    
    zneupd( &rvec, &howmany, select.data(),
            eigVal.data(), eigv, &ldv, &sigma, 
            workev.data(),  &bmat, &n, which.c_str(), &nev, &tol, 
            resid.data(), &ncv, eigv, &ldv, iparam, ipntr, workd.data(), 
            workl.data(), &lworkl, rwork.data(), &ierr );
  }

};

// template<class T, class EigVec_t>
// bool 
// ArpackSymSolverOMP<T,EigVec_t>::put(xmlDocPtr doc, xmlNsPtr ns, xmlNodePtr cur) {

//   xmlNodePtr tcur = cur->xmlChildrenNode;

//   while (tcur != NULL) {
//     if((!xmlStrcmp(tcur->name, (const xmlChar *)"set"))) {
//       string aname =  (const char*)(xmlGetProp(tcur, (const xmlChar *) "name"));
//       if(aname == "maxiter") {
// 	putXMLNodeContent<int>::put(maxiter,doc,tcur);
// 	iparam[2] = maxiter;
//       } else if(aname== "spectrum") {
// 	putXMLNodeContent<string>::put(which,doc,tcur);
//       } else if(aname == "tolerance") {
// 	putXMLNodeContent<real_type>::put(tol,doc,tcur);
//       } else if(aname == "arnoldivectors") {
// 	putXMLNodeContent<int>::put(ncv,doc,tcur);
//       }
//     }
//     tcur = tcur->next;
//   }
//   return true;
// }
#endif

/***************************************************************************
 * $RCSfile: ArpackSymSolverOMP.h,v $   $Author: jnkim $
 * $Revision: 1.1.1.1 $   $Date: 2004/07/23 21:04:02 $
 * $Id: ArpackSymSolverOMP.h,v 1.1.1.1 2004/07/23 21:04:02 jnkim Exp $ 
 ***************************************************************************/

